# encoding: utf-8
import sys
import cv2
import numpy
import numpy.matlib
import rospy
import time
from cv_bridge import CvBridge, CvBridgeError
from darknet_ros_msgs.msg import BoundingBoxes
from darknet_ros_msgs.msg import BoundingBox
from std_msgs.msg import String
from sensor_msgs.msg import Image
#from __future__ import print_function
from random import randint
#from skimage import img_as_float

class Tracking:
    def __init__(self, vehicle_type, vehicle_id):
        self.vehicle_type = vehicle_type
        self.vehicle_id = vehicle_id
        self.target_id = []  # 初始化id框
        self.target_class = []
        self.target_probability = []
        self.target_data = []
        self.target = BoundingBoxes()
        self.file=open('result_track'+str(self.vehicle_id)+'.txt','w') 

    def createTrackerByName(self, trackerType):
        trackerTypes = ['BOOSTING', 'MIL', 'KCF', 'TLD', 'MEDIANFLOW', 'GOTURN', 'MOSSE', 'CSRT']
        if trackerType == trackerTypes[0]:
            tracker = cv2.TrackerBoosting_create()
        elif trackerType == trackerTypes[1]:
            tracker = cv2.TrackerMIL_create()
        elif trackerType == trackerTypes[2]:
            tracker = cv2.TrackerKCF_create()
        elif trackerType == trackerTypes[3]:
            tracker = cv2.TrackerTLD_create()
        elif trackerType == trackerTypes[4]:
            tracker = cv2.TrackerMedianFlow_create()
        elif trackerType == trackerTypes[5]:
            tracker = cv2.TrackerGOTURN_create()
        elif trackerType == trackerTypes[6]:
            tracker = cv2.TrackerMOSSE_create()
        elif trackerType == trackerTypes[7]:
            tracker = cv2.TrackerCSRT_create()
        else:
            tracker = None
            print('Incorrect tracker name')
            print('Available trackers are:')
            for t in trackerTypes:
                print(t)
        return tracker

    def raw_pic_callback(self, data):
        global raw_pic
        a = CvBridge()
        try:
            raw_pic = a.imgmsg_to_cv2(data, "bgr8")
        except CvBridgeError as e:
            print(e)

    def pix_posi_callback(self, data):
        global predict_bbox, flag, predict_bbox_mad
        id_all = [0, 1, 2, 3, 4, 5]  # 标注的行人
        red_man0_id = 0  # 红色行人1
        red_man1_id = 5  # 红色行人2
        for target1 in data.bounding_boxes:
            if target1.id in id_all:
                pix_posi[0, target1.id] = (target1.xmax + target1.xmin) / 2
                pix_posi[1, target1.id] = (target1.ymax + target1.ymin) / 2
                self.target_id.append(target1.id)
                self.target_class.append(target1.Class)
                self.target_probability.append(target1.probability)
                predict_bbox_mid = list(predict_bbox)  # 将元组转换成列表
                bbox_mad = list(predict_bbox_mad)
                predict_bbox_mid[0] = float(target1.xmin)  # 存储数据
                predict_bbox_mid[1] = float(target1.ymin)
                predict_bbox_mid[2] = float(target1.xmax - target1.xmin)
                predict_bbox_mid[3] = float(target1.ymax - target1.ymin)
                predict_bbox = list(predict_bbox_mid)  # 转换回列表predict_bbox数据格式为[xmin,ymin,l,h]
                bbox_mad[0] = float(target1.probability)# 存储old数据
                bbox_mad[1] = int(target1.xmin)
                bbox_mad[2] = int(target1.ymin)
                bbox_mad[3] = int(target1.xmax)
                bbox_mad[4] = int(target1.ymax)
                bbox_mad[5] = int(target1.id)
                bbox_mad[6] = str(target1.Class)
                self.target_data.append(bbox_mad)
                flag = True
        self.target = data
        self.file.write('bounding_boxes   '+str(data.bounding_boxes))
        self.file.write('\n')
        #return self.target_id, self.target_class,self.target_probability,self.target_data,self.target

    def start(self):
        global  raw_pic,predict_bbox, pix_posi, flag, predict_bbox_mad,count,frame
        count = 0
        trackerTypes = ['BOOSTING', 'MIL', 'KCF', 'TLD', 'MEDIANFLOW', 'GOTURN', 'MOSSE', 'CSRT']
        predict_bbox = [0.0, 0.0, 0.0, 0.0]  # 初始化列表predict_bbox
        predict_bbox_mad = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
        pix_posi = numpy.matlib.zeros((2, 6))  # 初始化列表2*6，用于存储6架飞机的bbox的中心坐标
        flag = False
	
        # 订阅话题
        self.cmd_bboxes = rospy.Subscriber("uav_" + str(self.vehicle_id) + "/darknet_ros/bounding_boxes", BoundingBoxes, self.pix_posi_callback)
        self.cmd_frame = rospy.Subscriber('/' + str(self.vehicle_type) + '_' + str(self.vehicle_id) + '/cgo3_camera/image_raw', Image,self.raw_pic_callback)
        # 发布话题
        self.pub_track = rospy.Publisher("uav_" + str(self.vehicle_id) + "/darknet_ros/bounding_boxes_track",BoundingBoxes, queue_size=10)
        rospy.init_node(self.vehicle_type + '_' + self.vehicle_id + "_tracking")
        rate = rospy.Rate(10)
        trackerType ='KCF'
        bbox = predict_bbox
        frame = raw_pic
        while not rospy.is_shutdown():  # 当节点未关闭时
            bboxes_last = []
            bboxes = tuple(bboxes_last)
            colors = []
            frame = raw_pic  # 读取新的一帧图像
            bbox = predict_bbox  # 更新新的bbox
            bboxes_last.append(bbox)
            colors.append((randint(64, 255), randint(64, 255), randint(64, 255)))
            multiTracker = cv2.MultiTracker_create()
            timer = cv2.getTickCount()  # 开始计时
            fps = cv2.getTickFrequency() / (cv2.getTickCount() - timer)  # 计算帧率FPS
            for bbox in bboxes_last:
                multiTracker.add(self.createTrackerByName(trackerType), frame, tuple(bbox))
            ok, bboxes = multiTracker.update(frame)  # 开始跟踪
            if flag: #节点未关闭，目标检测有效  
                self.pub_track.publish(self.target)  # 发布旧的话题
                print('detection:', flag)
                cv2.putText(frame, "Detection succeed ", (100, 80), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (50, 170, 50), 2)
                flag = False
                count = 0
            else: #节点关闭，目标检测失效
                print('tracking:', ok)
                if ok:  # 跟踪成功
                #print(self.target_probability)
                    while count <= 5:
                        msg_new = BoundingBoxes()  # 修改要发布的话题的数据
                        new_bbox =BoundingBox()
                        for i, newbox in enumerate(bboxes):
                            new_bbox.probability = float(self.target_probability[i])
                            new_bbox. xmin= int(newbox[0])
                            new_bbox.ymin= int(newbox[1])
                            new_bbox.xmax= int(newbox[0] + newbox[2])
                            new_bbox.ymax = int(newbox[1] + newbox[3])
                            new_bbox.id = int(self.target_id[i])
                            new_bbox.Class = str(self.target_class[i])
                            msg_new.bounding_boxes.append(new_bbox)
                            p1 = (int(newbox[0]), int(newbox[1]))
                            p2 = (int(newbox[0] + newbox[2]), int(newbox[1] + newbox[3]))
                            #cv2.rectangle(frame, p1, p2, colors[i], 2, 1)
                        msg_new.image_header = self.target.image_header
                        self.pub_track.publish(msg_new)  # 发布新的话题
                        self.file.write('bounding_boxes_track   '+str(msg_new.bounding_boxes))
                        self.file.write('\n')
                        # cv2.putText(frame, "Tracking succeed ", (100, 80), cv2.FONT_HERSHEY_SIMPLEX, 0.75,
                        #             (50, 170, 50), 2)
                        count += 1
                    self.target_class = []
                    self.target_probability = []
                    self.target_id = []
                        #print(count)
                
                else: 
                    count = 0 
                    # 跟踪失败
                    # cv2.putText(frame, "Tracking failured ", (100, 80), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2)
                    # if count<=5:
                    #self.pub_track.publish(self.target)  # 发布旧的话题
            # cv2.putText(frame, trackerType + " Tracker", (100, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (50, 170, 50), 2)
            # cv2.putText(frame, "FPS : " + str(int(fps)), (100, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (50, 170, 50), 2)
            # cv2.imshow('MultiTracker', frame)
                    if cv2.waitKey(1) & 0xFF == 27:
                        break
        rate.sleep()


if __name__ == '__main__':
    tracking = Tracking(sys.argv[1], sys.argv[2])
    tracking.start()

